# -*- coding: utf-8 -*-
"""
Created on Mon Sep 28 16:57:26 2020

@author: Robin
"""

import numpy as np
import pandas as pd
from .performance import factor_information_coefficient,get_indicator_analysis
from .utils import (cal_factor_data,merge_to_multiindex)
from .utils import compute_forward_returns
#
#from lib.performance import factor_information_coefficient,get_indicator_analysis
#from lib.utils import (cal_factor_data,merge_to_multiindex)
#from lib.utils import compute_forward_returns


from sklearn.linear_model import LinearRegression
from sklearn.decomposition import PCA

def get_combined_factor_df(factors,stock,combine_type = 'ic_ir'):
    """
    传入因子数据、股票数据、
    计算日期、回溯时长
    ----------------------------------------------------
    model_type是因子合成的方法
    目前可选'ic_ir','pca'和'max_icir'
    'max_icir'法目前是没有约束条件下的解法，比较粗糙，有待完善
    -----------------------------------------------------
    factors是self.factors的数据，是一个双索引的DataFrame，索引为日期和股票代码，列名为因子名
                          indi1_in_day  indi2_in_day
    date       asset                                
    2020-09-10 000008.SZ      1.836478      1.810613
               000009.SZ      0.501466      1.552083
    
    2020-09-22 603328.SH     -0.919232      1.129487
               603338.SH     -1.057981     -0.850450           
    ------------------------------------------------------
    stock是self.stock的DataFrame，索引是日期，列名为股票代码
                 000008.SZ  000009.SZ   000012.SZ  ...  603888.SH  603939.SH  603983.SH
    2020-09-10  67.673045  60.038643  140.051163  ...  61.867993  279.44919  65.306056
    2020-09-11  67.448962  61.548490  143.532358  ...  62.393411  288.53064  66.541657
    2020-09-14  67.448962  63.680040  144.067926  ...  60.107842  287.34861  66.029334      
    --------------------------------------------------------
    """
    
    if combine_type=='ic_ir':
        '''ic_ir加权法合并因子'''
        weights = {}
        # 计算因子权重
        for ind_code in factors.columns:
            # 生成对齐数据
            factor_data = cal_factor_data(stock,factors[ind_code])
            # 计算ic-ir作为权重
            ictmp = factor_information_coefficient(factor_data[{'1D','factor'}],group_adjust=False,by_group=False)
            ic_ir = (ictmp.mean() / ictmp.std()).values[0]   # 越大越好
            weights[ind_code] = ic_ir
               
        # 因子合成
        combined_factor_df = sum(factors[ind_code] * weights[ind_code] for ind_code in factors.columns)
    
    elif combine_type=='max_icir':
        '''
        最大化ic_ir法加权合并因子
        '''
        IC_all=pd.DataFrame()
        IC_mean=[]
        
        for ind_code in factors.columns:
            factor_data = cal_factor_data(stock,factors[ind_code])
            ictmp = factor_information_coefficient(factor_data[{'1D','factor'}],group_adjust=False,by_group=False)
            IC_all[ind_code]=ictmp['1D']
            IC_mean.append(ictmp.mean().values[0])
        
        IC_all=IC_all.set_index(ictmp.index)
        IC_cov=IC_all.cov()
        
        #转化成计算需要的形式
        COV=IC_cov.as_matrix()
        IC=np.array(IC_mean).reshape(-1,1)
        
        #计算权重
        weight_array=np.dot(np.linalg.inv(COV),IC)
        
        '''
        注意：这里使用了最优化问题没有约束条件的显式解，
        后续应该使用带约束条件（权重>0）的二次规划解
        '''
        
        #合成字典
        weight_list=weight_array.reshape(1,-1).tolist()[0]
        weights=dict(zip(factors.columns,weight_list))
        # 因子合成
        combined_factor_df = sum(factors[ind_code] * weights[ind_code] for ind_code in factors.columns)
 
    
    elif combine_type =='pca':
        '''PCA主成分分析法合并因子'''
        date_index=factors.index.levels[0]
        stock_index=factors.index.levels[1]
        
        combined_factor_df=pd.DataFrame()
        
        for date in date_index:
            X=np.array(factors.loc[date].fillna(method='ffill'))
            #简单地用前一行的值填充缺失的因子值。后续应做更精细的处理。
            pca=PCA(n_components=1)
            pca.fit(X)
            X_reduced=pca.transform(X)
            
            combined_indi=pd.DataFrame(X_reduced,index=stock_index,columns=['combined_indi'])
            combined_indi['date']=date
            combined_factor_df=combined_factor_df.append(combined_indi)
    
        combined_factor_df=combined_factor_df.set_index([combined_factor_df['date'], combined_factor_df.index],drop=True)   
        combined_factor_df=combined_factor_df.drop(columns=['date'])
        combined_factor_df=combined_factor_df['combined_indi']
        #combined_factor_df是一个Series
    else:
        pass        
        
    return combined_factor_df
    '''
    函数get_combined_factor_df的返回值是一个双索引的DataFrame，索引为日期和股票代码
    内容是多个因子合成的一个因子的数据
        date        asset    
    2020-09-10  000008.SZ   -2.241617
                000009.SZ   -0.612092
                
    2020-09-22  603328.SH    1.122020
                603338.SH    1.291378
    '''


def get_top_stocks(df,stocknum,ic_mean):
    if ic_mean>=0:
        return df.sort_values(ascending = False)[:stocknum].index.values   
    else:
        return df.sort_values(ascending = True)[:stocknum].index.values
        
def get_stock_list_by_score(factor,groupby,datestr,stock_num,ic_mean):
    '''
    打分法，根据单个日期进行选股
    '''
    factor_name = factor.name            
    date_dt = pd.to_datetime(datestr)           
    dateindex = factor.index.levels[0]        
    select_date = dateindex[dateindex<=date_dt][-1]
    
    if (groupby is not None):
        # 分组
        factor = merge_to_multiindex(factor,groupby)
        select_data = factor.loc[select_date]     #series可以不用loc，但是dataframe必须用loc
        stock_df = select_data.groupby('industry')[factor_name].apply(get_top_stocks,stock_num,ic_mean)
        # 返回结果
        stock_choosed = stock_df.to_dict()
        stock_choosed_type = 2      
    else:   
        # 不分组
        select_data = factor.loc[select_date]     #series可以不用loc，但是dataframe必须用loc        
        
        # 排序筛选
        '''若IC均值为正，选因子值最大的股票；若IC均值为负，选因子值最小的股票'''
        if ic_mean>=0:
            stock_list = select_data.sort_values(ascending = False)[:stock_num].index.values  
        else:
             stock_list = select_data.sort_values(ascending = True)[:stock_num].index.values
      
        # 返回结果
        stock_choosed = list(stock_list)   
        stock_choosed_type = 1
    return stock_choosed,stock_choosed_type

def get_stock_list_by_score_by_dates(factors,stock,groupby,datelist,stock_num,lookbackPeriod,combine_type):
    '''
    打分法，根据多个日期进行选股
    -----------------------------------------------------------------------
    输入值factors是一个双索引的DataFrame，索引为日期和股票代码，列名为因子名
                          indi1_in_day  indi2_in_day
    date       asset                                
    2020-09-10 000008.SZ      1.836478      1.810613
               000009.SZ      0.501466      1.552083
               
    2020-09-22 603328.SH     -0.919232      1.129487
               603338.SH     -1.057981     -0.
    --------------------------------------------------------------------------
    stock是储存股票收盘价信息的DataFrame，索引为时间，列名为股票名
                000008.SZ  000009.SZ   000012.SZ  ...  603888.SH  603939.SH  603983.SH
    2020-09-10  67.673045  60.038643  140.051163  ...  61.867993  279.44919  65.306056
    2020-09-11  67.448962  61.548490  143.532358  ...  62.393411  288.53064  66.541657
    ---------------------------------------------------------------------------
    groupby是储存个股行业分类信息的双索引Series，索引为日期和股票代码
    date        asset    
    2007-01-31  000006.SZ     房地产
                000016.SZ    可选消费
                000028.SZ    医疗保健
                ....          ....
    2020-08-31  603328.SH    信息技术
                603338.SH      工业
                603355.SH    可选消费
    ---------------------------------------------------------------------------
    datelist是一个list,储存选股的日期;stocknum是选股的数量，int；
    combine_type是因子合成的方式，str，详见get_combined_factor_df的说明
    ---------------------------------------------------------------------------
    
    '''
    stock_choosed = {}  
    if (groupby is not None):
        # 分组        
        for datetmp in datelist:
            # 1.根据回望期以及日期序列合并相关因子，得到合并后的因子
            datetmp1 = pd.to_datetime(datetmp) # datetmp[:4]+'-'+datetmp[4:6]+'-'+datetmp[6:]
            # 提取相应日期的数据
            date_series_tmp = stock.index[stock.index<=datetmp1][-lookbackPeriod:]
            factorstmp = factors.loc[date_series_tmp]
            stocktmp = stock.loc[date_series_tmp]                
            # 因子合成
            combined_factor_df = get_combined_factor_df(factorstmp,stocktmp,combine_type)
            combined_factor_df.name = 'combine_factor'
            factor_name = 'combine_factor'
            select_date = date_series_tmp[-1]
            
            #获取IC均值
            combine_factor_data=cal_factor_data(stock,combined_factor_df)
            indicator_ana=get_indicator_analysis(combine_factor_data)
            ic_mean=indicator_ana.ic_analysis.ic_stats['IC_mean'][0]
            
            # 分组
            factor = merge_to_multiindex(combined_factor_df,groupby)
            select_data = factor.loc[select_date]     #series可以不用loc，但是dataframe必须用loc
            # 排序筛选
            stock_df = select_data.groupby('industry')[factor_name].apply(get_top_stocks,stock_num,ic_mean)
            # 选股结果
            stock_list = stock_df.to_dict()   
            stock_choosed[datetmp] = stock_list    
            
        stock_choosed_type = 4
    else:
        # 不分组
        for datetmp in datelist:
            '''
            # 按照factor里面的行业信息进行选股，设置一个是否需要使用行业信息来选股的参数
            '''
            # 1.根据回望期以及日期序列合并相关因子，得到合并后的因子
            datetmp1 = pd.to_datetime(datetmp) # datetmp[:4]+'-'+datetmp[4:6]+'-'+datetmp[6:]
            # 提取相应日期的数据
            date_series_tmp = stock.index[stock.index<=datetmp1][-lookbackPeriod:]
            factorstmp = factors.loc[date_series_tmp]
            stocktmp = stock.loc[date_series_tmp]
            
            # 因子合成
            combined_factor_df = get_combined_factor_df(factorstmp,stocktmp,combine_type)
            combined_factor_df.name = 'combine_factor'
            factor_name = 'combine_factor'
            
            #获取IC均值
            combine_factor_data=cal_factor_data(stock,combined_factor_df)
            indicator_ana=get_indicator_analysis(combine_factor_data)
            ic_mean=indicator_ana.ic_analysis.ic_stats['IC_mean'][0]
            
            select_date = date_series_tmp[-1] 
            # 不分组
            select_data = combined_factor_df.loc[select_date]     #series可以不用loc，但是dataframe必须用loc        
            # 排序筛选
            if ic_mean>=0:
                stock_list = select_data.sort_values(ascending = False)[:stock_num].index.values  
            else:
                stock_list = select_data.sort_values(ascending = True)[:stock_num].index.values
            
            # 返回选股结果
            stock_list = list(stock_list) 
            stock_choosed[datetmp] = stock_list     
                   
        stock_choosed_type = 3
        
    return stock_choosed,stock_choosed_type
    '''
    返回值stock_choosed_type代表选股的类型，3代表不行业分类，4代表行业分类
    如果行业分类了，返回值stock_choosed是一个字典，键是日期，值是一个字典
    值的字典中键是行业，值是选股的list；
    如果不行业分类，返回值stock_choosed是一个字典，键是日期，值是选股的list
    '''

def get_stock_list_by_regression(factors,stock,groupby,datestr,stock_num):
    '''
    -----------------------------------------------------------------------
    输入值factors是一个双索引的DataFrame，索引为日期和股票代码，列名为因子名
                          indi1_in_day  indi2_in_day
    date       asset                                
    2020-09-10 000008.SZ      1.836478      1.810613
               000009.SZ      0.501466      1.552083
               
    2020-09-22 603328.SH     -0.919232      1.129487
               603338.SH     -1.057981     -0.
    --------------------------------------------------------------------------
    stock是储存股票收盘价信息的DataFrame，索引为时间，列名为股票名
                000008.SZ  000009.SZ   000012.SZ  ...  603888.SH  603939.SH  603983.SH
    2020-09-10  67.673045  60.038643  140.051163  ...  61.867993  279.44919  65.306056
    2020-09-11  67.448962  61.548490  143.532358  ...  62.393411  288.53064  66.541657
    ---------------------------------------------------------------------------
    groupby是储存个股行业分类信息的双索引Series，索引为日期和股票代码
    date        asset    
    2007-01-31  000006.SZ     房地产
                000016.SZ    可选消费
                000028.SZ    医疗保健
                ....          ....
    2020-08-31  603328.SH    信息技术
                603338.SH      工业
                603355.SH    可选消费
    ---------------------------------------------------------------------------
    【注意】这个选股方法暂时不支持行业分类的功能！需要后续补充

    '''
    
    if (groupby is not None):
        # 分组
        stock_choosed_type = 2      
    else:     
        stock_choosed_type = 1
    
    select_ind = list(factors.columns)
    forward_returns = compute_forward_returns(factors[select_ind[0]],stock)
    
    date_dt = pd.to_datetime(datestr)
    dateindex = factors.index.levels[0]        
    select_date = dateindex[dateindex<=date_dt][-1]
    
    #用day1的收益率数据作为因变量进行回归
    ret_series=forward_returns.loc[select_date]['1D']        
    #用day0的因子数据作为自变量进行回归
    indi_date=dateindex[dateindex<=date_dt][-2]        
    indi_df=factors.loc[indi_date]
    
    #多元线性回归
    Y=ret_series.fillna(0)
    X=indi_df.fillna(0)
    
    lr = LinearRegression()
    lr.fit(X,Y)
    
    #用day1的因子值预测day2的收益率数据
    X_use=factors.loc[select_date]
    Y_pred = lr.predict(X_use)
    
    predict_ret=pd.Series(Y_pred,index=stock.columns)
    stock_list = predict_ret.sort_values(ascending = False)[:stock_num].index.values                 
            
    # 返回结果
    stock_choosed = list(stock_list) 

    return stock_choosed,stock_choosed_type

def get_stock_list_by_regress_by_dates(factors,stock,groupby,datelist,stock_num,lookbackPeriod):
    '''
    注意：使用回归法选股函数的时候，factors和stock的最后时间必须超过datelist的最后时间
    因为要使用前一天的数据对后一天的收益率进行回归，所以如果输入的股票价格数据仅仅到选股当天，
    那么无法计算当天的收益率，则无法预测下一期的收益率，即无法选股
    '''
    
    stock_choosed = {}
    
    if (groupby is not None):
        # 分组
        stock_choosed_type = 4   
    else:     
        stock_choosed_type = 3
    
    select_ind = list(factors.columns)
    forward_returns = compute_forward_returns(factors[select_ind[0]],stock)
    dateindex = factors.index.levels[0]
    
    if datelist[-1]>=dateindex[-1]:
        #选股的日期超过了能预测的日期
        print('提供的收益率数据不足，无法预测',datelist[-1].strftime('%Y-%m-%d'),'的收益率，无法选股')
        return None,stock_choosed_type
    
    for datetmp in datelist:
        select_date = dateindex[dateindex<=datetmp][-1]
    
        #用day1的收益率数据作为因变量进行回归
        ret_series=forward_returns.loc[select_date]['1D']        
        #用day0的因子数据作为自变量进行回归
        indi_date=dateindex[dateindex<=datetmp][-2]        
        indi_df=factors.loc[indi_date]
    
        #多元线性回归
        Y=ret_series.fillna(0)
        X=indi_df.fillna(0)
    
        lr = LinearRegression()
        lr.fit(X,Y)
        
        #用day1的因子值预测day2的收益率数据
        X_use=factors.loc[select_date]
        Y_pred = lr.predict(X_use)
        
        predict_ret=pd.Series(Y_pred,index=stock.columns)
        stock_list = predict_ret.sort_values(ascending = False)[:stock_num].index.values                 
                    
        # 返回结果
        stock_choosed[datetmp] = list(stock_list) 
    
    return stock_choosed,stock_choosed_type
